from fastapi import Depends
from sqlmodel import select, update, desc, delete, SQLModel

from typing_extensions import deprecated
import typing as T
from pymilvus import MilvusClient
from ...init.Global import Session, EmbeddingSession
from ...pojo import UserKnowledgeBase, UserKnowledgeBaseVO
from .BaseMapper import BaseMapper

# 用户知识库的相关操作的Mapper[对外]


class _UserKnowledgeBaseMapper(BaseMapper):

    embeddingSession: MilvusClient = None

    def __call__(self, session: Session, embeddingSession: EmbeddingSession):
        self.session = session
        self.embeddingSession = embeddingSession
        return self

    async def delete_knowledgebase_by_id(
        self, knowledgeBaseId: int, collection_name: str = "default"
    ):
        await self.session.execute(
            delete(UserKnowledgeBase).where(UserKnowledgeBase.id == knowledgeBaseId)
        )
        self.embeddingSession.delete(
            collection_name, filter=f"knowledgeBaseId == {knowledgeBaseId}"
        )

    # 将知识插入知识库
    async def vector_insert_batch(
        self, data: T.List[T.Dict[str, T.Any]], collection_name: str = "default"
    ):
        self.embeddingSession.insert(collection_name, data)

    # 根据表达式和嵌入向量进行搜索
    async def vector_search(
        self,
        query_embeddiing: list[float],
        exp: str,
        top_n: int = 1,
        collection_name: str = "default",
        nprobe: int = 16,
    ):
        result = (
            self.embeddingSession.search(
                collection_name,
                [query_embeddiing],
                filter=exp,
                output_fields=["content", "knowledgeBaseId", "metadata", "userId"],
                limit=top_n,
                anns_field="embedding",
                search_params={
                    "params": {
                        "nprobe": nprobe,
                    }
                },
            )
        )[0]
        return result

    async def vector_query_by_expression(
        self, exp: str, count: int, collection_name: str = "default"
    ):
        if count == -1:
            result = self.embeddingSession.query(
                collection_name=collection_name,
                filter=exp,
                output_fields=["content", "metadata", "userId", "knowledgeBaseId"],
            )
        else:
            result = self.embeddingSession.query(
                collection_name=collection_name,
                filter=exp,
                output_fields=["content", "metadata", "userId", "knowledgeBaseId"],
                limit=count,
            )
        return result

    # 根据id删除知识
    async def vector_delete_by_id(self, id: int, collection_name: str = "default"):
        self.embeddingSession.delete(collection_name, ids=[id])


UserKnowledgeBaseMapper = T.Annotated[
    _UserKnowledgeBaseMapper, Depends(_UserKnowledgeBaseMapper())
]
